import sys  
import cv2  
import mediapipe as mp 
from PyQt5.QtCore import QTimer  
from PyQt5.QtGui import QImage, QPixmap  
from PyQt5.QtWidgets import QApplication, QMainWindow  
from new_ui import Ui_MainWindow  
import numpy as np 
import os 
import time 
import tempfile 

from 模型测验 import predict_pull_up, mp_pose, mp_drawing


class MainApp(QMainWindow):  
    def __init__(self):  
        super().__init__() 
        self.ui = Ui_MainWindow() 
        self.ui.setupUi(self)  

        
        self.pose = mp_pose.Pose(
            static_image_mode=False,
            model_complexity=1, 
            min_detection_confidence=0.5, 
            min_tracking_confidence=0.5  
        )
        
        self.cap = None 
        self.timer = QTimer() 
        self.timer.timeout.connect(self.update_frame)  
        
        # 引体向上检测相关变量
        self.detection_timer = QTimer() 
        self.detection_timer.timeout.connect(self.detect_pull_up) 
        self.temp_dir = tempfile.mkdtemp() 
        self.success_count = 0  
        self.last_prediction = 0 
        self.last_success_time = 0 
        self.cooldown_period = 1.5  # 冷却时间（秒）用于防止误判

       
        self.ui.close_camera_btn.setEnabled(False) 
        self.ui.open_camera_btn.clicked.connect(self.open_camera)  
        self.ui.close_camera_btn.clicked.connect(self.close_camera) 

    def open_camera(self): 
        """打开摄像头并开始视频捕获"""
        self.cap = cv2.VideoCapture(0)  
        if self.cap.isOpened():  
            self.ui.close_camera_btn.setEnabled(True)  
            self.ui.open_camera_btn.setEnabled(False)  
            self.timer.start(30) 
            self.detection_timer.start(300) 
            self.ui.text_cos.setText("引体向上检测已启动，当前成功次数: 0")  
        else:
            print("Error: 无法访问摄像头")  

    def close_camera(self):  
        """停止视频捕获并释放摄像头"""
        self.timer.stop()  
        self.detection_timer.stop() 
        if self.cap:  
            self.cap.release()  
        self.ui.open_camera_btn.setEnabled(True)  
        self.ui.close_camera_btn.setEnabled(False)  
        self.clear_video_labels()  
       
        self.success_count = 0
        self.last_prediction = 0

    def update_frame(self):  
        """捕获并处理视频帧"""
        ret, frame = self.cap.read() 
        if not ret: 
            return  
        self.display_frame(frame, self.ui.original_video_label)

        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) 
        results = self.pose.process(rgb_frame) 

        processed_frame = frame.copy()
        if results.pose_landmarks:  
            mp_drawing.draw_landmarks(
                processed_frame, 
                results.pose_landmarks,  
                mp_pose.POSE_CONNECTIONS, 
                mp_drawing.DrawingSpec(color=(0,255,0), thickness=2, circle_radius=2), 
                mp_drawing.DrawingSpec(color=(255,0,0), thickness=2)  
            )

        self.display_frame(processed_frame, self.ui.processed_video_label)

    def detect_pull_up(self):
        """截取当前帧并进行引体向上检测"""
        if not self.cap or not self.cap.isOpened():
            return
            
        ret, frame = self.cap.read()
        if not ret:
            return
            
        temp_image_path = os.path.join(self.temp_dir, f"frame_{int(time.time()*1000)}.jpg")
        cv2.imwrite(temp_image_path, frame)
        
        prediction, probability = predict_pull_up(temp_image_path)
        
        try:
            os.remove(temp_image_path)
        except:
            pass
            
        if prediction is None:
            return
        
        current_time = time.time()
        if prediction == 1:  

            if (self.last_prediction == 0 and
                (current_time - self.last_success_time) > self.cooldown_period):
                self.success_count += 1
                self.last_success_time = current_time
                self.ui.text_cos.setText(f"引体向上检测 - 成功次数: {self.success_count} (概率: {probability:.2f})")
        
        self.last_prediction = prediction

    def display_frame(self, frame, target_label):  
        """将OpenCV图像帧显示到PyQt标签控件"""
        if frame is None or target_label is None:  
            return

        h, w, ch = frame.shape 
        bytes_per_line = ch * w  
        qt_image = QImage(frame.data, w, h, bytes_per_line, QImage.Format_BGR888)  
        pixmap = QPixmap.fromImage(qt_image).scaled(  
            target_label.size(), aspectRatioMode=1  
        )
        target_label.setPixmap(pixmap) 

    def clear_video_labels(self):
        """清除视频显示区域的内容"""
        self.ui.original_video_label.clear() 
        self.ui.original_video_label.setText("摄像头已关闭") 
        self.ui.processed_video_label.clear()  
        self.ui.processed_video_label.setText("摄像头已关闭") 
        self.ui.text_cos.setText("")  

    def closeEvent(self, event): 
        self.close_camera() 
        try:
            import shutil
            shutil.rmtree(self.temp_dir, ignore_errors=True)
        except:
            pass
        super().closeEvent(event)  

if __name__ == "__main__": 
    app = QApplication(sys.argv)  
    main_window = MainApp()  
    main_window.show()  
    sys.exit(app.exec_()) 